import numpy as np
import matplotlib.pyplot as plt


# EXERCISE 1
# To test your solution run ex1_parallel_coordinates.py
def parallel_coordinates(coordinates, values, labels, cmap):
  """Plot 2d array `values` using K parallel coordinates.

  Arguments:

    coordinates -- list or array of K elements containg coordinate
      names,
    values -- (K,N)-shaped array of N data points with K
      coordinates,
    labels -- list or array of one string per data point
      describing its class membership (category)
  """

  # SOLUTION
  ax = plt.subplot(111)

  # find names and number of different classes
  ulabels = np.unique(labels)
  n_labels = len(ulabels)

  # for each select distinct colors from Accent pallette
  colors = cmap(np.arange(n_labels) * cmap.N / (n_labels + 1))

  # change the label strings to indices into class names array
  class_id = np.searchsorted(ulabels, labels)
  lines = plt.plot(values[:, :], 'k')
  [l.set_color(colors[c]) for c, l in zip(class_id, lines)]

  # add grid, configure labels and axes
  ax.spines['top'].set_visible(False)
  ax.spines['bottom'].set_position(('outward', 5))
  ax.spines['bottom'].set_visible(False)
  ax.yaxis.set_ticks_position('both')
  ax.xaxis.set_ticks_position('none')

  plt.xticks(np.arange(len(coordinates)), coordinates)
  plt.grid(axis='x', ls='-')

  leg_handlers = [lines[np.where(class_id == id)[0][0]] for id in range(n_labels)]
  ax.legend(leg_handlers,
            ulabels,
            frameon=False,
            loc='upper left',
            ncol=len(labels),
            bbox_to_anchor=(0, -0.03, 1, 0))

  # END OF SOLUTION


# EXERCISE 2
# To test your solution run ex2_summer_olympics.py
def format_data(data, sel_disciplines, sel_countries, sel_years):
  """Count the medals won by selected countries in selected
  discplines over selected Summer Olympics edition.
  Shape data in format suitable for parallel_coordinates.

  Arguments:

    data -- record array imported from CSV. Each row in the array
       contains data about a single medal-wining athlete. Fields
       relevant to this function are: `noc` (country codes),
       `edition` (year of the olympics), `discipline`

    sel_disciplines -- numpy array of selected discipline names

    sel_countries -- numpy array of selected country codes

    sel_years -- numpy array of selected Olympic years as integers

  Returns:

    medal_count -- 2D array of total number of medals aggregated
       over all athletes in selected
       disciplines (columns) grouped by Olympic years and country of
       origin (rows). Each row represent one country and one year,
       but different countries can be intermixed.

    country_names -- 1D numpy array with one entry (string) per row
       of `medal_counts`, which identifies the row with country
       name.
  """
  medal_year = data['edition']
  medal_country = data['noc']
  medal_discipline = data['discipline']

  # SOLUTION
  medal_count = []
  country_names = []
  for noc in sel_countries:
    for yr in sel_years:
      # group data by year/country
      country_year_discipline = medal_discipline[(medal_country == noc) & (medal_year == yr)]
      # count medals in selected disciplines
      medals = [(country_year_discipline == disc).sum() for disc in sel_disciplines]
      medal_count.append(medals)
      country_names.append(noc)
  medal_count = np.array(medal_count).T
  # END OF SOLUTION

  medal_count = np.asarray(medal_count)
  country_names = np.asarray(country_names)

  return medal_count, country_names


# EXERCISE 3
# To test your solution run ex3_normalised_coordinates.py
def normalised_coordinates(coordinates, values, labels, cmap):
  """Plot 2d array `values` using K parallel coordinates.

  Arguments:

    coordinates -- list or array of K elements containg coordinate
      names,
    values -- (K,N)-shaped array of N data points with K
      coordinates,
    labels -- list or array of one string per data point
      describing its class membership (category)
  """

  K = len(coordinates)

  # normalise data
  vmin, vmax = values.min(0), values.max(0)
  normed_values = (values - vmin) * 1. / (vmax - vmin + 1e-9)

  # SOLUTION

  # create independent, attached axes
  fig, axes = plt.subplots(1, K - 1)
  fig.subplots_adjust(wspace=0)

  # find names and number of different classes
  ulabels = np.unique(labels)
  n_labels = len(ulabels)

  # obtain colors
  colors = cmap(np.arange(n_labels) * cmap.N / (n_labels + 1))

  # change the label strings to indices into class names array
  class_id = np.searchsorted(ulabels, labels)

  for i in range(K - 1):
    ax = axes[i]
    lines = ax.plot(normed_values, 'k')

    # set line colors
    [l.set_color(colors[c]) for c, l in zip(class_id, lines)]

    # configure axes
    ax.spines['top'].set_visible(False)
    ax.spines['bottom'].set_position(('outward', 5))
    ax.spines['bottom'].set_visible(False)
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('none')

    # set limit to show only single line segment
    ax.set_xlim((i, i + 1))
    ax.set_xticks([i])
    ax.set_xticklabels([coordinates[i]])

    # set the scale
    ax.set_yticks([0, 1])
    ax.set_yticklabels([vmin[i], vmax[i]])

  # we have to deal with rightmost axis separately
  axes[-1].set_xticks([K - 2, K - 1])
  axes[-1].set_xticklabels(coordinates[-2:])

  ax = axes[-1].twinx()
  ax.tick_params(axis='y', direction='out')
  ax.xaxis.set_ticks_position('none')
  ax.set_yticks([0, 1])
  ax.set_yticklabels([vmin[-1], vmax[-1]])

  # END OF SOLUTION

  # add legend in figure coordinates
  leg_handlers = [lines[np.where(class_id == idx)[0][0]] for idx in range(len(coordinates))]
  plt.legend(leg_handlers,
             ulabels,
             frameon=False,
             loc='upper left',
             ncol=len(labels),
             bbox_to_anchor=(0.1, 0.06, 0.9, 0.02),
             bbox_transform=fig.transFigure)

  # set leftmost axis as the default for labelling
  plt.sca(axes[0])


def normalised_coordinates2(coordinates, values, labels, cmap, title):
  import matplotlib
  matplotlib.use("Agg")
  import matplotlib.pyplot as plt
  from matplotlib.path import Path
  import matplotlib.patches as patches
  import numpy as np

  ynames = coordinates
  ys = values
  ymins = ys.min(axis=0).astype(float) * (1 - 1e-4)
  ymaxs = ys.max(axis=0).astype(float) * (1 + 1e-4)
  dys = ymaxs - ymins

  # transform all data to be compatible with the main axis
  zs = np.zeros_like(ys)
  zs[:, 0] = ys[:, 0]
  zs[:, 1:] = (ys[:, 1:] - ymins[1:]) / dys[1:] * dys[0] + ymins[0]

  fig, host = plt.subplots()
  y_formatter = matplotlib.ticker.ScalarFormatter(useOffset=False)

  axes = [host] + [host.twinx() for i in range(ys.shape[1] - 1)]
  for i, ax in enumerate(axes):
    if ymins[i] != ymaxs[i]:
      ax.set_ylim(ymins[i], ymaxs[i])
    ax.spines['top'].set_visible(False)
    ax.spines['bottom'].set_visible(False)
    if ax != host:
      ax.spines['left'].set_visible(False)
      ax.yaxis.set_ticks_position('right')
      ax.spines["right"].set_position(("axes", i / (ys.shape[1] - 1)))
    ax.yaxis.set_major_formatter(y_formatter)

  if ys.shape[1] > 1:
    host.set_xlim(0, ys.shape[1] - 1)
  host.set_xticks(range(ys.shape[1]))
  host.set_xticklabels(ynames, fontsize=6)
  host.tick_params(axis='x', which='major', pad=7)
  host.spines['right'].set_visible(False)
  host.xaxis.tick_top()
  host.yaxis.set_major_formatter(y_formatter)

  labels_unique = np.unique(labels).tolist()
  nlabelus = len(labels_unique)

  legend_handles = [None for _ in labels_unique]
  for j in range(ys.shape[0]):
    ncol = ys.shape[1]
    # create bezier curves
    verts = list(
        zip([x for x in np.linspace(0, ncol - 1, ncol * 3 - 2, endpoint=True)],
            np.repeat(zs[j, :], 3)[1:-1]))
    codes = [Path.MOVETO] + [Path.CURVE4 for _ in range(len(verts) - 1)]
    path = Path(verts, codes)
    juniq = labels_unique.index(labels[j])
    patch = patches.PathPatch(path, facecolor='none', lw=2, edgecolor=cmap(juniq / nlabelus))
    legend_handles[juniq] = patch
    host.add_patch(patch)
  host.set_title(title, fontsize=8, pad=12)

  plt.subplots_adjust(left=0.1, right=0.9, top=0.8, bottom=0.2)

  host.legend(legend_handles, ["%.2f" % label for label in labels_unique],
              loc='lower center',
              bbox_to_anchor=(0.5, -0.18),
              ncol=len(labels),
              fontsize=6)
